import xarray as xr
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from matplotlib.pyplot import pcolormesh

# 读取数据
ds = xr.open_dataset('./wrf_data/wrfout_d02_2024-05-01_12:00:00')
cldfra = ds['CLDFRA'].isel(Time=0)  # 选择第一个时间步

# 正确获取经纬度坐标(处理可能的维度差异)
if 'Time' in ds['XLAT'].dims:
    lats = ds['XLAT'].isel(Time=0).values  # 如果XLAT有时间维度
    lons = ds['XLONG'].isel(Time=0).values
else:
    lats = ds['XLAT'].values
    lons = ds['XLONG'].values

# 验证维度
print(f"CLDFRA维度: {cldfra.dims}")
print(f"纬度数组维度: {lats.shape}")
print(f"经度数组维度: {lons.shape}")

#print(lats)
#print(lons)

# 确保纬度/经度与CLDFRA的水平维度匹配
assert lats.shape == (cldfra.sizes['south_north'], cldfra.sizes['west_east'])
assert lons.shape == (cldfra.sizes['south_north'], cldfra.sizes['west_east'])

#print(cldfra)
#exit()

nlev = cldfra.sizes['bottom_top']
nlat = cldfra.sizes['south_north']
nlon = cldfra.sizes['west_east']

cld_locations = np.zeros((nlat,nlon))
print(nlev,nlat,nlon)
for j in range(nlat):
    for i in range(nlon):
        if np.max(cldfra[:,j,i].values) > 0.:
            #cld_locations[j,i] = np.max(cldfra[:,j,i].values)
            cld_locations[j,i] = 1.

#  # 找出有云的网格点
#  cloud_locations = []
#  for k in range(cldfra.sizes['bottom_top']):
#      layer_data = cldfra.isel(bottom_top=k)
#      #print(layer_data[0,:].values)
#      #exit()
#      cloud_positions = np.where(layer_data > 0)  # 获取有云的行列索引
#
#      for i, j in zip(*cloud_positions):
#          try:
#              cloud_locations.append({
#                  'latitude': lats[i, j],
#                  'longitude': lons[i, j],
#                  'vertical_layer': k,
#                  'cloud_fraction': float(layer_data[i, j])})
#          except IndexError as e:
#              print(f"索引错误发生在层{k}, 位置({i},{j})")
#              print(f"纬度数组形状: {lats.shape}")
#              print(f"经度数组形状: {lons.shape}")
#              raise
#
#  # 转换为DataFrame
#  if cloud_locations:  # 如果有云才创建DataFrame
#      df_clouds = pd.DataFrame(cloud_locations)
#      print("\n有云的网格点信息:")
#      print(df_clouds)
#  else:
#      print("未检测到云覆盖")

fig, ax = plt.subplots()
#plt.scatter(df_clouds['longitude'], df_clouds['latitude'], c=df_clouds['vertical_layer'], s=5)
im = ax.pcolormesh(lons, lats, cld_locations)
fig.colorbar(im, ax=ax,label='cloud fraction')
plt.title('Cloud Fraction')
#plt.savefig('/mnt/raid1/lijing/WRF_model/output/cldfra_0.png')
plt.show()
